#include "cybertron/core/UtilLz4.hpp"

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>


#include <lz4.h>
#include <lz4frame.h>

CYBERTRON_BEGIN

CybertronCoreAPI int UtilLz4::compressBound(int inputSize) {
	return LZ4_compressBound(inputSize);
}

CybertronCoreAPI int UtilLz4::compressData(const char* src, char* dst, int srcSize, int dstCapacity) {
	return LZ4_compress_default(src, dst, srcSize, dstCapacity);
}

CybertronCoreAPI int UtilLz4::decompressData(const char* src, char* dst, int compressedSize, int dstCapacity) {
	return LZ4_decompress_safe(src, dst, compressedSize, dstCapacity);
}


#define IN_CHUNK_SIZE  (16*1024)

static const LZ4F_preferences_t kPrefs = {
	{ LZ4F_max256KB, LZ4F_blockLinked, LZ4F_noContentChecksum, LZ4F_frame,
	0 /* unknown content size */, 0 /* no dictID */ , LZ4F_noBlockChecksum },
	0,   /* compression level; 0 == default */
	0,   /* autoflush */
	0,   /* favor decompression speed */
	{ 0, 0, 0 },  /* reserved, must be set to 0 */
};

/* safe_fwrite() :
* performs fwrite(), ensure operation success */
static bool safe_fwrite(const char* buf, size_t size, FILE* f)
{
	if (size == 0) return true;

	size_t const writtenSize = fwrite(buf, 1, size, f);

	if (writtenSize < size) {
		if (ferror(f))  /* note : ferror() must follow fwrite */ 
			return false;
		else {
			return safe_fwrite(buf + writtenSize, size - writtenSize, f);
		}
	}
	return true;
}

typedef struct {
	int error;
	unsigned long long size_in;
	unsigned long long size_out;
} compressResult_t;

static compressResult_t
compress_file_internal(FILE* f_in, FILE* f_out,
	LZ4F_compressionContext_t ctx,
	char* inBuff, size_t inChunkSize,
	char* outBuff, size_t outCapacity)
{
	compressResult_t result = { 1, 0, 0 };  /* result for an error */
	unsigned long long count_in = 0, count_out;

	assert(f_in != NULL); assert(f_out != NULL);
	assert(ctx != NULL);
	assert(outCapacity >= LZ4F_HEADER_SIZE_MAX);
	assert(outCapacity >= LZ4F_compressBound(inChunkSize, &kPrefs));

	/* write frame header */
	size_t const headerSize = LZ4F_compressBegin(ctx, outBuff, outCapacity, &kPrefs);
	if (LZ4F_isError(headerSize)) {
		return result;
	}
	count_out = headerSize;
	safe_fwrite((char*)outBuff,  headerSize, f_out);

	/* stream file */
	for (;;) {
		size_t const readSize = fread(inBuff, 1, IN_CHUNK_SIZE, f_in);
		if (readSize == 0) break; /* nothing left to read from input file */
		count_in += readSize;

		size_t const compressedSize = LZ4F_compressUpdate(ctx,
			outBuff, outCapacity,
			inBuff, readSize,
			NULL);

		if (LZ4F_isError(compressedSize)) {
			return result;
		}

		if (compressedSize > 0)
			safe_fwrite(outBuff, compressedSize, f_out);

		count_out += compressedSize;
	}

	/* flush whatever remains within internal buffers */
	size_t const compressedSize = LZ4F_compressEnd(ctx,
		outBuff, outCapacity,
		NULL);
	if (LZ4F_isError(compressedSize)) 
		return result;

	safe_fwrite(outBuff, compressedSize, f_out);

	count_out += compressedSize;

	result.size_in = count_in;
	result.size_out = count_out;
	result.error = 0;
	return result;
}

static compressResult_t
compress_file(FILE* f_in, FILE* f_out)
{
	assert(f_in != NULL);
	assert(f_out != NULL);

	/* ressource allocation */
	LZ4F_compressionContext_t ctx;
	size_t const ctxCreation = LZ4F_createCompressionContext(&ctx, LZ4F_VERSION);
	char* src = (char*)malloc(IN_CHUNK_SIZE);
	size_t const outbufCapacity = LZ4F_compressBound(IN_CHUNK_SIZE, &kPrefs);   /* large enough for any input <= IN_CHUNK_SIZE */
	char* outbuff = (char*)malloc(outbufCapacity);

	compressResult_t result = { 1, 0, 0 };  /* == error (default) */
	if (!LZ4F_isError(ctxCreation) && src && outbuff) {
		result = compress_file_internal(f_in, f_out,
			ctx,
			src, IN_CHUNK_SIZE,
			outbuff, outbufCapacity);
	}

	LZ4F_freeCompressionContext(ctx);   /* supports free on NULL */
	free(src);
	free(outbuff);
	return result;
}

static size_t get_block_size(const LZ4F_frameInfo_t* info) {
	switch (info->blockSizeID) {
	case LZ4F_default:
	case LZ4F_max64KB:  return 1 << 16;
	case LZ4F_max256KB: return 1 << 18;
	case LZ4F_max1MB:   return 1 << 20;
	case LZ4F_max4MB:   return 1 << 22;
	default:
		// "Impossible with expected frame specification (<=v1.6.1)\n";
		assert(false );
	}
	return 1 << 16;
}

/* @return : 1==error, 0==success */
static int
decompress_file_internal(FILE* f_in, FILE* f_out,
	LZ4F_dctx* dctx,
	void* src, size_t srcCapacity, size_t filled, size_t alreadyConsumed,
	void* dst, size_t dstCapacity)
{
	int firstChunk = 1;
	size_t ret = 1;

	assert(f_in != NULL); assert(f_out != NULL);
	assert(dctx != NULL);
	assert(src != NULL); assert(srcCapacity > 0); assert(filled <= srcCapacity); assert(alreadyConsumed <= filled);
	assert(dst != NULL); assert(dstCapacity > 0);

	/* Decompression */
	while (ret != 0) {
		/* Load more input */
		size_t readSize = firstChunk ? filled : fread(src, 1, srcCapacity, f_in); firstChunk = 0;
		const void* srcPtr = (const char*)src + alreadyConsumed; alreadyConsumed = 0;
		const void* const srcEnd = (const char*)srcPtr + readSize;
		if (readSize == 0 || ferror(f_in)) {
			// printf("Decompress: not enough input or error reading file\n");
			return 1;
		}

		/* Decompress:
		* Continue while there is more input to read (srcPtr != srcEnd)
		* and the frame isn't over (ret != 0)
		*/
		while (srcPtr < srcEnd && ret != 0) {
			/* Any data within dst has been flushed at this stage */
			size_t dstSize = dstCapacity;
			size_t srcSize = (const char*)srcEnd - (const char*)srcPtr;
			ret = LZ4F_decompress(dctx, dst, &dstSize, srcPtr, &srcSize, /* LZ4F_decompressOptions_t */ NULL);
			if (LZ4F_isError(ret)) {
				// printf("Decompression error: %s\n", LZ4F_getErrorName(ret));
				return 1;
			}
			/* Flush output */
			if (dstSize != 0) safe_fwrite((char*)dst, dstSize, f_out);
			/* Update input */
			srcPtr = (const char*)srcPtr + srcSize;
		}

		assert(srcPtr <= srcEnd);

		/* Ensure all input data has been consumed.
		* It is valid to have multiple frames in the same file,
		* but this example only supports one frame.
		*/
		if (srcPtr < srcEnd) {
			printf("Decompress: Trailing data left in file after frame\n");
			return 1;
		}
	}

	/* Check that there isn't trailing data in the file after the frame.
	* It is valid to have multiple frames in the same file,
	* but this example only supports one frame.
	*/
  
	size_t const readSize = fread(src, 1, 1, f_in);
	if (readSize != 0 || !feof(f_in)) {
		// printf("Decompress: Trailing data left in file after frame\n");
		return 1;
	}  

	return 0;
}


/* @return : 1==error, 0==completed */
static int
decompress_file_allocDst(FILE* f_in, FILE* f_out,
	LZ4F_dctx* dctx,
	void* src, size_t srcCapacity)
{
	assert(f_in != NULL); assert(f_out != NULL);
	assert(dctx != NULL);
	assert(src != NULL);
	assert(srcCapacity >= LZ4F_HEADER_SIZE_MAX);  /* ensure LZ4F_getFrameInfo() can read enough data */

												  /* Read Frame header */
	size_t const readSize = fread(src, 1, srcCapacity, f_in);
	if (readSize == 0 || ferror(f_in)) {
		// printf("Decompress: not enough input or error reading file\n");
		return 1;
	}

	LZ4F_frameInfo_t info;
	size_t consumedSize = readSize;
	   
	size_t const fires = LZ4F_getFrameInfo(dctx, &info, src, &consumedSize);
	if (LZ4F_isError(fires)) {
		// printf("LZ4F_getFrameInfo error: %s\n", LZ4F_getErrorName(fires));
		return 1;
	}   

	/* Allocating enough space for an entire block isn't necessary for
	* correctness, but it allows some memcpy's to be elided.
	*/
	size_t const dstCapacity = get_block_size(&info);
	void* const dst = malloc(dstCapacity);
	if (!dst) { 
		// perror("decompress_file(dst)"); 
		return 1; 
	}

	int const decompressionResult = decompress_file_internal(
		f_in, f_out,
		dctx,
		src, srcCapacity, readSize - consumedSize, consumedSize,
		dst, dstCapacity);

	free(dst);
	return decompressionResult;
}


/* @result : 1==error, 0==success */
static int decompress_file(FILE* f_in, FILE* f_out)
{
	assert(f_in != NULL); assert(f_out != NULL);

	/* Ressource allocation */
	void* const src = malloc(IN_CHUNK_SIZE);
	if (!src) { 
		// perror("decompress_file(src)"); 
		return 1; 
	}

	LZ4F_dctx* dctx;
 
	LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION);

	int const result = !dctx ? 1 /* error */ :
		decompress_file_allocDst(f_in, f_out, dctx, src, IN_CHUNK_SIZE);

	free(src);
	LZ4F_freeDecompressionContext(dctx);   /* note : free works on NULL */
	return result;
}

CybertronCoreAPI int UtilLz4::compressFile(const char* srcFile, const char* dstFile) {
	FILE* const inpFp = fopen(srcFile, "rb");
	FILE* const outFp = fopen(dstFile, "wb");

	compressResult_t const ret = compress_file(inpFp, outFp);

	fclose(outFp);
	fclose(inpFp);

	return ret.error;
}

CybertronCoreAPI int UtilLz4::decompressFile(const char* srcFile, const char* dstFile) {
	FILE* const inpFp = fopen(srcFile, "rb");
	FILE* const outFp = fopen(dstFile, "wb");

	int const ret = decompress_file(inpFp, outFp);

	fclose(outFp);
	fclose(inpFp);

	return ret;
}

//////////////////////////////////////////////////////////////////////////

/* @return : 1==error, 0==success */
static int
decompress_file_internal_mem(FILE* f_in, UtilLz4MemWriter& dst,
	LZ4F_dctx* dctx,
	void* src, size_t srcCapacity, size_t filled, size_t alreadyConsumed,
	size_t dstCapacity)
{
	int firstChunk = 1;
	size_t ret = 1;

	assert(f_in != NULL);
	assert(dctx != NULL);
	assert(src != NULL); assert(srcCapacity > 0); assert(filled <= srcCapacity); assert(alreadyConsumed <= filled);
	assert(dstCapacity > 0);

	/* Decompression */
	while (ret != 0) {
		/* Load more input */
		size_t readSize = firstChunk ? filled : fread(src, 1, srcCapacity, f_in); firstChunk = 0;
		const void* srcPtr = (const char*)src + alreadyConsumed; alreadyConsumed = 0;
		const void* const srcEnd = (const char*)srcPtr + readSize;
		if (readSize == 0 || ferror(f_in)) {
			// printf("Decompress: not enough input or error reading file\n");
			return 1;
		}

		/* Decompress:
		* Continue while there is more input to read (srcPtr != srcEnd)
		* and the frame isn't over (ret != 0)
		*/
		while (srcPtr < srcEnd && ret != 0) {
			/* Any data within dst has been flushed at this stage */
			char* dstbuff = dst.alloc(dstCapacity);
			size_t dstSize = dstCapacity;
			size_t srcSize = (const char*)srcEnd - (const char*)srcPtr;

			ret = LZ4F_decompress(dctx, dstbuff, &dstSize, srcPtr, &srcSize, /* LZ4F_decompressOptions_t */ NULL);
			if (LZ4F_isError(ret)) {
				// printf("Decompression error: %s\n", LZ4F_getErrorName(ret));
				return 1;
			}
			/* Flush output */
			if (dstSize != 0) dst.write(dstSize);
			/* Update input */
			srcPtr = (const char*)srcPtr + srcSize;
		}

		assert(srcPtr <= srcEnd);

		/* Ensure all input data has been consumed.
		* It is valid to have multiple frames in the same file,
		* but this example only supports one frame.
		*/
		if (srcPtr < srcEnd) {
			// printf("Decompress: Trailing data left in file after frame\n");
			return 1;
		}
	}

	/* Check that there isn't trailing data in the file after the frame.
	* It is valid to have multiple frames in the same file,
	* but this example only supports one frame.
	*/

	size_t const readSize = fread(src, 1, 1, f_in);
	if (readSize != 0 || !feof(f_in)) {
		// printf("Decompress: Trailing data left in file after frame\n");
		return 1;
	}

	return 0;
}

/* @return : 1==error, 0==completed */
static int
decompress_file_allocDst_mem(FILE* f_in, UtilLz4MemWriter& dst,
	LZ4F_dctx* dctx,
	void* src, size_t srcCapacity)
{
	assert(f_in != NULL);
	assert(dctx != NULL);
	assert(src != NULL);
	assert(srcCapacity >= LZ4F_HEADER_SIZE_MAX);  /* ensure LZ4F_getFrameInfo() can read enough data */

												  /* Read Frame header */
	size_t const readSize = fread(src, 1, srcCapacity, f_in);
	if (readSize == 0 || ferror(f_in)) {
		return 1;
	}

	LZ4F_frameInfo_t info;
	size_t consumedSize = readSize;

	size_t const fires = LZ4F_getFrameInfo(dctx, &info, src, &consumedSize);
	if (LZ4F_isError(fires)) {
		return 1;
	}

	/* Allocating enough space for an entire block isn't necessary for
	* correctness, but it allows some memcpy's to be elided.
	*/
	size_t const dstCapacity = get_block_size(&info);

	if (consumedSize > readSize - consumedSize)
	{
		return 1;
	}

	int const decompressionResult = decompress_file_internal_mem(
		f_in, dst,
		dctx,
		src, srcCapacity, readSize - consumedSize, consumedSize,
		dstCapacity);


	return decompressionResult;
}

/* @result : 1==error, 0==success */
static int decompress_file_mem(FILE* f_in, UtilLz4MemWriter& dst)
{
	assert(f_in != NULL);

	/* Ressource allocation */
	void* const src = malloc(IN_CHUNK_SIZE);
	if (!src) {
		// perror("decompress_file(src)"); 
		return 1;
	}

	LZ4F_dctx* dctx;

	LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION);

	int const result = !dctx ? 1 /* error */ :
		decompress_file_allocDst_mem(f_in, dst, dctx, src, IN_CHUNK_SIZE);

	free(src);
	LZ4F_freeDecompressionContext(dctx);   /* note : free works on NULL */
	return result;
}
CybertronCoreAPI int UtilLz4::decompressFileToMem(const char* srcFile, UtilLz4MemWriter& dst) {
	FILE* const inpFp = fopen(srcFile, "rb");

	int const ret = decompress_file_mem(inpFp, dst);

	fclose(inpFp);

	return ret;
}
CYBERTRON_END